import os
import cv2 as cv
import numpy as np
import torch.nn as nn
from torch.autograd import Variable
import torch
from sklearn.model_selection import train_test_split
import torchvision.models as tvmodel

device = torch.device('cuda')

def readData(path):
    x_images = []
    y_labels = []
    image_size = 100
    for i, j in enumerate(os.listdir(path)):
        sub_path = os.path.join(path, j)
        for file_name in os.listdir(sub_path):
            file_path = os.path.join(sub_path, file_name)
            img = cv.imread(file_path)
            img = cv.resize(img, dsize=(image_size, image_size)) / 255.
            x_images.append(img)
            y_labels.append(i)
    return np.array(x_images), np.array(y_labels)


DATA_PATH = 'zoo'
x_images, y_labels = readData(DATA_PATH)
x_images = Variable(torch.Tensor(x_images))
x_images = torch.transpose(x_images, 1, 3)
y_labels = Variable(torch.Tensor(y_labels))

x_train, x_test, y_train, y_test = train_test_split(x_images, y_labels, train_size=0.8, shuffle=True)

class MyMODEL(nn.Module):
    def __init__(self, num_classes=1000):
        super(MyMODEL, self).__init__()
        restnet = tvmodel.resnet34(pretrained=True)
        restnet_out_channels = restnet.fc.in_features
        self.resnet = nn.Sequential(*list(restnet.children())[:-2])
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.layers = nn.Linear(restnet_out_channels, num_classes)

    def forward(self, inputs):
        x = self.resnet(inputs)
        x = self.avgpool(x)
        x = torch.squeeze(x)
        x = self.layers(x)
        return x


if __name__ == '__main__':
    model = MyMODEL(num_classes=2)
    losses = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(params=model.parameters(), lr=0.001)
    epochs = 10
    for epoch in range(epochs):
        optimizer.zero_grad()
        pred = model(x_train)
        loss = losses(pred, y_train.long())
        loss.backward()
        optimizer.step()

        acc_train = (torch.argmax(pred, 1) == y_train).float().mean()
        acc_test = (torch.argmax(model(x_test), 1) == y_test).float().mean()
        print('epoch:', epoch + 1, 'loss_val:', loss.item(), '训练集准确率:', acc_train.detach().numpy(), '测试集准确率:',acc_test.detach().numpy())